import SimpleITK as sitk
import numpy as np
import os

from mycode import section_synthesis


secData_root = 'mydata_section'

sparse_root = os.path.join(secData_root, 'sparse_landmark')
synthesis_root = os.path.join(secData_root, 'synthesis_landmark')
section_root = os.path.join(secData_root, 'section')

pred_sparse_root = os.path.join(secData_root, 'pred_sparse_landmark')
pred_synthesis_root = os.path.join(secData_root, 'pred_synthesis_landmark')
pred_section_root = os.path.join(secData_root, 'pred_section')

sax_gt_root = os.path.join('mydata_gt', 'rotate', 'output_shortAxis')
sax_pred_root = os.path.join('mydata_pred', 'rotate', 'output_shortAxis')

SPARSE_PREFIX = 'sparse_'
SYNTHESIS_PREFIX = 'synthesis_'
NIFTI_SUFFIX = '.nii.gz'
SAX_PREFIX = 'case_'
SEG_NAME = 'segmentation.nii.gz'
SECTION_PREFIX = 'section_'

def segment_attach(seg_0: np.ndarray, seg_1: np.ndarray):
    pass


def gen_section(seq_list, is_pred=False):
    """

    :param seq_list:
    :param is_pred: whether the segment and landmark is predicted
    :return:
    """
    for seq_number in seq_list:

        if is_pred:
            sax_root = sax_pred_root
            landmark_path = os.path.join(pred_synthesis_root, SYNTHESIS_PREFIX + seq_number + NIFTI_SUFFIX)
            section_path = os.path.join(pred_section_root, SECTION_PREFIX + seq_number + NIFTI_SUFFIX)
        else:
            sax_root = sax_gt_root
            landmark_path = os.path.join(synthesis_root, SYNTHESIS_PREFIX + seq_number + NIFTI_SUFFIX)
            section_path = os.path.join(section_root, SECTION_PREFIX + seq_number + NIFTI_SUFFIX)

        segment_path = os.path.join(sax_root, SAX_PREFIX + seq_number, SEG_NAME)
        seg_img = sitk.ReadImage(segment_path)
        section_array = section_synthesis.get_section(sitk.GetArrayFromImage(seg_img).astype(np.uint8), sitk.GetArrayFromImage(sitk.ReadImage(landmark_path)))
        section_img = sitk.GetImageFromArray(section_array)
        section_img.CopyInformation(seg_img)
        sitk.WriteImage(section_img, section_path)
        print('finished section:', seq_number)


def gen_synthesis_landmark(seq_list, is_pred=False):
    for seq_number in seq_list:
        if is_pred:
            sparse_path = os.path.join(pred_sparse_root, SPARSE_PREFIX + seq_number + NIFTI_SUFFIX)
            synthesis_path = os.path.join(pred_synthesis_root, SYNTHESIS_PREFIX + seq_number + NIFTI_SUFFIX)
        else:
            sparse_path = os.path.join(sparse_root, SPARSE_PREFIX + seq_number + NIFTI_SUFFIX)
            synthesis_path = os.path.join(synthesis_root, SYNTHESIS_PREFIX + seq_number + NIFTI_SUFFIX)
        if not os.path.exists(sparse_path):
            continue
        section_synthesis.landmark_synthesis(sparse_path, synthesis_path)
        print('finished synthesis of', seq_number)

# not tested yet
if __name__ == '__main__':
    train_list = [144, 146, 147, 149, 152, 154, 155, 158, 159, 160, 161, 162, 166, 
					167, 168, 170, 171, 172, 173, 176]
    train_list += list(range(200, 258))
    
    ori_test_list = list(range(177, 197))
    ori_test_list.remove(186)

    extra_test_list = list(range(260, 345))
    extra_test_list += list(range(350, 442))

    seq_list = [str(i).zfill(5) for i in ori_test_list + extra_test_list]

    seq_list = ['00177', '00308', '00441']

    # gen_synthesis_landmark(seq_list, is_pred=False)
    gen_section(seq_list, is_pred=False)
